package graceful

// Golang 程序 优雅的重启与关闭
// 重启: kill -USR2 pid, kill -HUP pid
// 关闭: kill -INT pid, kill -QUIT pid
// 注: 直接使用run运行, 重启后临时文件会发生变化, 从而导致重启失败

import (
	"context"
	"errors"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"syscall"
	"time"
)

var (
	isRunChaning   bool // 是否正在运行协程通道程序
	isCheckChaning bool // 是否正在检查中
	chanSignal     chan int
	workIDArr      = []bool{}

	listener net.Listener
	server   *http.Server
	timeout  = time.Second * 30
	graceful = flag.Bool("graceful", false, "listen on fd open 3 (internal use only)")
)

// NewWorker 新建任务
func NewWorker() int {
	var index int = 0
	for k, v := range workIDArr {
		if !v {
			return k
		}
		index++
	}
	if len(workIDArr) < index+1 {
		workIDArr = append(workIDArr, true)
	}
	isRunChaning = true
	return index
}

// EndWork 结束任务
func EndWork(id int) {
	if len(workIDArr) > id {
		workIDArr[id] = false
		isRunChaning = getLiveWorkNumber() > 0
		if chanSignal != nil && isCheckChaning {
			chanSignal <- id
		}
	}
}

func getLiveWorkNumber() (ret int) {
	for _, v := range workIDArr {
		if v {
			ret++
		}
	}
	if ret == 0 {
		isRunChaning = false
	}
	return ret
}

// ListenAndServe ..
func ListenAndServe(addr string, handler http.Handler) {
	var err error
	server = &http.Server{
		Addr:    addr,
		Handler: handler,
	}
	flag.Parse()
	if *graceful {
		fmt.Println("net.FileListener startup")
		f := os.NewFile(3, "")
		defer f.Close()
		if listener, err = net.FileListener(f); err != nil {
			fmt.Println("net.FileListener:", err)
			return
		}
	} else {
		fmt.Println("net.Listen startup")
		if listener, err = net.Listen("tcp", server.Addr); err != nil {
			panic(err)
		}
	}
	// 启动服务
	go func() {
		log.Printf("Listening and serving HTTP [%s], PID [%d]\n", addr, os.Getpid())
		if err := server.Serve(listener); err != nil {
			fmt.Printf("server.Serve err: %v\n", err)
		}
	}()
	chanSignal = make(chan int)
	// 监听信号
	handleSignals()
}

func handleSignals() {
	sig := make(chan os.Signal)
	// signal.Notify(sig, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGTSTP, syscall.SIGQUIT)
	signal.Notify(sig)
	for {
		capturedSig := <-sig
		// fmt.Printf("Received SIG. [PID:%d, SIG:%v]\n", syscall.Getpid(), capturedSig)

		switch capturedSig {
		case syscall.SIGHUP, syscall.SIGUSR2:
			// 优雅重启
			fmt.Println("graceful reload start")
			if err := reload(); err != nil {
				fmt.Println("graceful reload error", err)
			}
			stop()
			fmt.Println("graceful reload end")
			return
		case syscall.SIGINT, syscall.SIGTERM, syscall.SIGTSTP, syscall.SIGQUIT:
			// 优雅关闭
			fmt.Println("graceful shutdown start")
			signal.Stop(sig)
			stop()
			fmt.Println("graceful shutdown end")
			return
		}
	}
}

func reload() error {
	tl, ok := listener.(*net.TCPListener)
	if !ok {
		return errors.New("listener is not tcp listener")
	}
	f, err := tl.File()
	if err != nil {
		return err
	}
	defer f.Close()
	args := []string{"-graceful"}
	cmd := exec.Command(os.Args[0], args...)
	cmd.Stdout = os.Stdout         // 标准输出
	cmd.Stderr = os.Stderr         // 错误输出
	cmd.ExtraFiles = []*os.File{f} // 文件描述符
	return cmd.Start()
}

// checkWaitWorker 检测是否还有正在进行的协程
func checkWaitWorker() {
	if !isRunChaning {
		return
	}
	waitSecond := 30
	tick := time.Tick(time.Second * 1)
	fmt.Printf("还有未完成的协程正在运行(等待超过%d秒将强制结束) ~~\n", waitSecond)
	isCheckChaning = true
	defer func() {
		isCheckChaning = false
	}()
	for {
		select {
		case <-chanSignal:
			fmt.Print(">> 有个协程已完成, ")
			n := getLiveWorkNumber()
			if n == 0 {
				fmt.Println("所有都完成 END!")
				return
			}
			fmt.Printf("还剩%d个还未完成\n", n)
		case <-tick:
			waitSecond--
			if waitSecond == 0 {
				fmt.Println("强制结束")
				return
			} else if waitSecond < 5 {
				fmt.Println("倒计时:", waitSecond)
			} else if n := waitSecond % 5; n == 0 {
				fmt.Println("倒计时:", waitSecond)
			}
		}
	}
}

func stop() {
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	server.Shutdown(ctx)
	checkWaitWorker()
}
